import angr
from angrutils import plot_func_graph
import random
from collections import deque
from .parsing_util import Parser
import string
from angr.knowledge_plugins.functions import Function


def is_hex(s):
    hex_digits = set(string.hexdigits)
    return all(c in hex_digits for c in s)


class AngrUtil:
    def __init__(self, file_path) -> None:
        self.p = angr.Project(file_path, load_options={'auto_load_libs': False})
        self.idfer = self.p.analyses.Identifier()
        self.cfg = self.p.analyses.CFGFast()

        pass

    def identify_function(self):
        """Return all the function names
        Those func names all contain "compute" string
        """
        function_names = []
        for funcInfo in self.idfer.func_info:
            if "compute" in funcInfo.name:
                function_names.append(funcInfo.name)
        self.function_names = function_names
        print(self.function_names)

    def get_func_info_by_name(self, func_name):
        """Get func info by the name of the function

        Args:
            func_name (str): The name of the function

        Returns:
            FuncInfo: A data structure about function info (name, addr, ...)
        """

        for funcInfo in self.idfer.func_info:
            if func_name == funcInfo.name:
                return funcInfo
        return None

    def print_all_funcs(self):
        for funcInfo in self.idfer.func_info:
            print(hex(funcInfo.addr), funcInfo.name)

    def print_func_by_name(self, name):
        blocks = self.get_func_blocks_by_name(name)
        for b in blocks:
            b.pp()

    def get_all_funcs(self):
        self.func_dict = {}
        for funcInfo in self.idfer.func_info:
            self.func_dict[hex(funcInfo.addr)] = funcInfo.name
        return self.func_dict

    def get_function_graph(self, func_name):
        """Return a graph for the corresponding function name

        Args:
            func_name (str): The name of function

        Returns:
            graph: A NetworkX graph class
        """
        graph = None
        for func in self.cfg.functions.values():
            if func.name == func_name:
                graph = func.transition_graph
                break
        if graph is not None:
            remove_nodes = []
            for n in graph.nodes:
                if n.size == 0:
                    remove_nodes.append(n)
            for n in remove_nodes:
                graph.remove_node(n)
        return graph

    def get_func_blocks_by_name(self, func_name):
        """Return basic blocks of the specific function

        Args:
            func_name (str): The name of function

        Returns:
            list: A list of Blocks (the angr Block class).
        """
        for func in self.cfg.functions.values():
            if func.name == func_name:
                return list(func.blocks)
        return None

    def contruct_block_dict(self, blocks):
        """Build up a dictionary for <addr, block>

        Args:
            blocks (Block): A list of Angr Block

        Returns:
            dict: A <addr, block> dict
        """
        blocks_dict = {}
        for b in blocks:
            if b.addr not in blocks_dict:
                blocks_dict[b.addr] = b
        return blocks_dict

    def random_walk(self, blocks_dict, graph, max_length=100):
        """Given a dict of blocks and a block graph (In that graph,
        all the node is BlockNode, we have to get Block by using the addr)

        Args:
            blocks_dict (dict): A <addr, block> dict
            graph (Graph): A NetworkX Graph class
            max_length (int): Maximum length of a trace

        Returns:
            list: A trace of Angr Blocks
        """
        first_block = list(graph.nodes)[0]
        trace = [first_block]
        while len(list(graph.successors(trace[-1]))) > 0 \
                and len(trace) < max_length:
            block_list = list(graph.successors(trace[-1]))
            # Remove blocks with 0 size
            # That can handle some special blocks
            # (e.g. Function UnresolvableCellTarget)
            new_list = []
            for b in block_list:
                if b.size > 0:
                    new_list.append(b)
            block_list = new_list

            next_block = block_list[random.randint(0, len(block_list) - 1)]
            trace.append(next_block)

        block_trace = []
        for t in trace:
            if t.addr in blocks_dict:
                block_trace.append(blocks_dict[t.addr])
        return block_trace

    def print_trace(self, trace):
        for t in trace:
            disasm_block = t.disassembly
            for ins in disasm_block.insns:
                print("{}\t{}".format(ins.mnemonic, ins.op_str))

    def parsing_trace(self, trace):

        tokens = []
        for t in trace:
            disasm_block = t.disassembly
            for ins in disasm_block.insns:
                tokens.append(ins.mnemonic)
                oprands = Parser.parse_one_line(ins.op_str)
                for opr in oprands:
                    if opr.isdigit() or "0x" in opr or is_hex(opr[:-1]):
                        tokens.append("num")
                    elif opr == ",":
                        continue
                    else:
                        tokens.append(opr)
        return tokens

    def get_block_tokens(self, block):
        tokens = []
        disasm_block = block.disassembly
        for ins in disasm_block.insns:
            tokens.append(ins.mnemonic)
            # print(ins.op_str)
            oprands = Parser.parse_one_line(ins.op_str)
            for opr in oprands:
                # print(opr)
                if opr.isdigit() or "0x" in opr or is_hex(opr[:-1]):
                    if opr.isdigit():
                        tokens.append(opr)
                    elif "0x" in opr:
                        if "00000000" in opr:
                            opr = opr.replace("00000000", "")

                        if abs(int(opr, 16)) <= 2**15:
                            tokens.append(opr)
                        else:
                            tokens.append("num")
                    elif is_hex(opr[:-1]) and len(opr[:-1]) > 0:
                        tmp = "0x{}".format(opr[:-1])
                        if abs(int(tmp, 16)) <= 2**15:
                            tokens.append(tmp)
                        else:
                            tokens.append("num")

                elif opr == ",":
                    continue
                else:
                    tokens.append(opr)
        return tokens

    def plot_function_graph(self, graph, name):
        """Plot a function graph and save as a file

        Args:
            graph (networkx): Function Graph
            name (str): The name of the saved file
        """
        plot_func_graph(self.p, graph, name, asminst=True, vexinst=False)

    def check_possible_implicit_func(self, tmp_block):
        if self.func_dict is None:
            self.get_all_funcs()

        disasm_block = tmp_block.disassembly
        for idx, ins in enumerate(disasm_block.insns):
            if "[" not in ins.op_str:
                continue
            if ins.mnemonic == "lea" and "rip" in ins.op_str \
                    and "rdi" in ins.op_str and "0x" in ins.op_str:
                addr = ins.op_str.split("[")[-1][:-1].split("+ ")[-1]
                addr = int(addr, 16)
                final_addr = hex(addr + disasm_block.insns[idx + 1].address)
                if final_addr in self.func_dict:
                    return self.func_dict[final_addr]
                # Search other possible
                basic = int(final_addr, 16)
                for i in range(-40, 41, 1):
                    tmp_addr = hex(basic + i)
                    print("Start searching the address [{}]".format(tmp_addr))
                    if tmp_addr in self.func_dict:
                        return self.func_dict[tmp_addr]

        return ""

    def extend_graph(self, graph, node, node_successors, insert_graph):
        # Remove edges
        for ns in node_successors:
            graph.remove_edge(node, ns)

        stack = deque()
        visited_addr = set()

        for (n, degree) in insert_graph.in_degree:
            if degree == 0:
                graph.add_edge(node, n)
                stack.appendleft(n)
                break

        new_block_dict_arr = []
        # Build graph
        while len(stack) > 0:
            tmp_node = stack.pop()
            if tmp_node.addr not in visited_addr:
                visited_addr.add(tmp_node.addr)
                if tmp_node.size == 0:
                    continue
                # if not hasattr(tmp_node, 'successors'):
                #     continue
                try:
                    for ns in tmp_node.successors():
                        #  Check the function node!
                        if isinstance(ns, Function):
                            print("Find Function!")
                            new_block_dict_arr.append(self.contruct_block_dict(ns.blocks))
                            tmp_graph = ns.transition_graph
                            graph, tmp_d_arr = self.extend_graph(graph, tmp_node, [], tmp_graph)
                            for tmp_d in tmp_d_arr:
                                new_block_dict_arr.append(tmp_d)
                            # Remove Function Node
                            insert_graph.remove_node(ns)
                            insert_graph.add_nodes_from(tmp_graph.nodes)

                        else:
                            graph.add_edge(tmp_node, ns)
                            stack.appendleft(ns)

                except ValueError:
                    continue

        for (n, degree) in insert_graph.out_degree:
            if degree == 0:
                for ns in node_successors:
                    if ns.size > 0:
                        graph.add_edge(n, ns)

        return graph, new_block_dict_arr

    def try_extend_graph(self, graph, blocks_dict):
        nodes = graph.nodes
        self.get_all_funcs()

        print("Before Extension, node number: {}".format(len(nodes)))
        has_new_node = True
        visited_node = set()
        while has_new_node:
            has_new_node = False
            for n in nodes:
                addr = n.addr
                tmp_block = None
                tmp_name = ""
                if addr in blocks_dict:
                    tmp_block = blocks_dict[addr]
                if tmp_block is not None:
                    tmp_name = self.check_possible_implicit_func(tmp_block)
                if tmp_name != "" and tmp_name not in visited_node:
                    break

            if tmp_name != "":
                has_new_node = True
                visited_node.add(tmp_name)
                print("Add new node [{}]".format(tmp_name))
                graph, new_block_dict_arr = self.extend_graph(graph, n, n.successors(),
                                                              self.get_function_graph(tmp_name))
                new_dict = self.contruct_block_dict(self.get_func_blocks_by_name(tmp_name))
                # Append blocks of the tmp_name function
                for k in new_dict:
                    blocks_dict[k] = new_dict[k]
                # Append blocks of extra function nodes
                for new_block_dict in new_block_dict_arr:
                    for k in new_block_dict:
                        blocks_dict[k] = new_block_dict[k]

        print("After Extension, node number: {}".format(len(graph.nodes)))
        return graph, blocks_dict


if __name__ == "__main__":
    file_path = "./tvm_data/other_funcs/elf/" \
        + "ZeroPadding2D-__padding___1,__input_shape____100,_100,_10__.so"
    util = AngrUtil(file_path)
    util.identify_function()
    g = util.get_function_graph(util.function_names[0])
    blocks = util.get_func_blocks_by_name(util.function_names[0])
    blocks_dict = util.contruct_block_dict(blocks)
    g, blocks_dict = util.try_extend_graph(g, blocks_dict)
    util.plot_function_graph(g, "cfg")

    for _ in range(3):
        trace = util.random_walk(blocks_dict, g)
        tokens = util.parsing_trace(trace)
        print(tokens)

    # util.print_trace(util.random_walk(blocks_dict, g))

    # print(blocks_dict.keys())

    # print(g.nodes)
    # first_block = list(g.nodes)[0]

    # disasm_block = blocks_dict[first_block.addr].disassembly
    # for ins in disasm_block.insns:
    #     print("{}\t{}".format(ins.mnemonic, ins.op_str))

    # for ins in blocks_dict[first_block.addr]:
    #     print(ins)

    # for b in g.successors(first_block):
    #     if b.size > 0:
    #         print(hex(b.addr))
    #         print(blocks_dict[b.addr].pp())
